# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

from .types import State
from .nodes import (
    coordinator_node,
    planner_node,
    reporter_node,
    research_team_node,
    researcher_node,
    coder_node,
    human_feedback_node,
    background_investigation_node,
)


def _build_base_graph():
    """构建并返回包含所有节点和边的基础状态图。"""
    builder = StateGraph(State)  # 使用State类型初始化状态图
    
    # 添加节点和边的连接关系
    builder.add_edge(START, "coordinator")  # 起始点连接到协调器（边）
    builder.add_node("coordinator", coordinator_node)  # 协调器节点 （节点）
    builder.add_node("background_investigator", background_investigation_node)  # 背景调查节点（节点）
    builder.add_node("planner", planner_node)  # 规划器节点（节点）
    builder.add_node("reporter", reporter_node)  # 报告生成节点（节点）
    builder.add_node("research_team", research_team_node)  # 研究团队节点（节点）
    builder.add_node("researcher", researcher_node)  # 研究员节点（节点）
    builder.add_node("coder", coder_node)  # 程序员节点（节点）
    builder.add_node("human_feedback", human_feedback_node)  # 人工反馈节点（节点）
    builder.add_edge("reporter", END)  # 报告节点连接到终点（边）
    return builder


# 在app.py中使用这个
def build_graph_with_memory():
    """构建并返回带有记忆功能的代理工作流图。"""
    # 使用持久化内存保存对话历史
    memory = MemorySaver()
    
    # 构建状态图并添加记忆功能
    builder = _build_base_graph()
    return builder.compile(checkpointer=memory)


def build_graph():
    """构建并返回不带记忆功能的代理工作流图。"""
    builder = _build_base_graph()
    return builder.compile()


graph = build_graph()
